from model import CNN_BE,UNet
from dataset import Trainset_PDE_generated,Trainset_noise_generated,Trainset_random_generated,Trainset_random_generated0


import time
import math
import os
import torch
import numpy as np
# from Numerical_Ju import numerical,mass_engy
from FFT import numerical_all
import torch.optim as optim
from torch import linalg
from torch.optim.lr_scheduler import StepLR
from loss import loss_be
from utils import save_checkpoints
from numerical_solution import Numerical_solution,operatorx,operatory,bioperatorx,bioperatory,bioperatorxy

import matplotlib.cm as cm
import matplotlib.pyplot as plt
torch.set_default_tensor_type(torch.DoubleTensor)

os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
class Tester():
    def __init__(self):

        self.height =256
        self.width =256
        self.intial_snap = 1
        self.mid_channels = 16
        self.deltaT = 0.05
        self.trainingBatch = 10
        self.testBatch = 10
        self.domain = [-0.5,0.5,-0.5,0.5]
        self.type = 'BE'
        self.epsilon = 0.01
        self.TotalSnaps = 1
        self.TotalTrainingBatch = 200

        self.hx = 1/self.height
        self.wx = 1/self.width

        self.device = 'cuda'
        self.num_type = 'CN'
        self.model_type = 'CNN'
        # self.model_type = 'UNet'

        if self.model_type == 'CNN':
            self.BECNN = CNN_BE(self.mid_channels)
        elif self.model_type == 'UNet':
            self.BECNN = UNet(self.intial_snap, 1)

        self.PDETrainSet =Trainset_random_generated0(self.testBatch,self.intial_snap, self.height, self.width)
        self.Snap_Training = self.PDETrainSet()

        self.PDETrainSet0 = Trainset_random_generated0(self.testBatch,self.intial_snap, self.height, self.width)
        self.Snap_Training0 = self.PDETrainSet0()

        # self.PDETrainSet1 = Trainset_PDE_generated(self.intial_snap, self.height, self.width, self.epsilon,self.domain)
        # self.Snap_Training = self.PDETrainSet1()
        # self.Snap_Training0 = self.PDETrainSet1()
        # self.Numerical_Full = np.load(
        #     f'data/deltaT{1e-2}_{self.height}_{self.width}_{888}_full.npy')
        # self.Pre = np.load(f'data/deltaT0.1_256_256_2R_16_888_full.npy')
        # self.Num = np.load(f'data/deltaT0.01_256_256_555_Numerical_all.npy')
        # self.Num=np.load(f'data/deltaT0.01_256_256_888_full.npy')
        # self.Num = np.load(f'data/deltaT0.01_256_256_666_Numerical_all.npy')
        print(self.BECNN)

        if self.device == 'cuda':
            self.BECNN = self.BECNN.to(self.device)
            self.Snap_TrainingLB = self.Snap_Training.to(self.device)
            self.Snap_TrainingLB0 = self.Snap_Training0.to(self.device)
            # self.Snap_TrainingLB = self.Snap_Training.to(self.device)[:, :, 1:, 0:-1]
            # self.Snap_TrainingLB0 = self.Snap_Training0.to(self.device)[:, :, 1:, 0:-1]


        self.model_name = f'NS_{self.model_type}_deltaT{self.deltaT}_{self.trainingBatch}_{self.width}_{self.height}_{self.type}_{self.TotalTrainingBatch}Batch{self.TotalSnaps}Snaps_{self.mid_channels}Channels'
    def full(self,miss):

        full_left = torch.clone(miss[:, :, :, [0]])
        full = torch.cat((miss, full_left), 3)
        full_bottom = torch.clone(full[:, :, [-1], :])
        full = torch.cat((full_bottom, full), 2)

        return full

    def test(self,TimeSteps):
        self.BECNN.eval()
        best_model = torch.load(f'checkpoints/{self.model_name}/checkpoint.pth.tar')
        self.BECNN.load_state_dict(best_model['state_dict'])
        Snap_init = torch.clone(self.Snap_TrainingLB0)
        Snap_init0 = torch.clone(self.Snap_TrainingLB0)
        # Snap_init0 = torch.clone(self.Snap_TrainingLB0)[0][0].cpu().detach().numpy()
        initial= self.full(Snap_init)[0][0].cpu().detach().numpy()
        # np.save(f'data/deltaT{self.deltaT}_{self.height}_{self.width}_{88}_initial.npy',initial)


        TimeSteps =1
        n = TimeSteps
        #
        # # a,b =mass_engy(self.full(Snap_init)[0][0].cpu().detach().numpy(),self.hx,self.wx,self.epsilon)
        # # c= np.max(np.abs(self.full(Snap_init)[0][0].cpu().detach().numpy()))
        # # max=[c]
        # # mass = [a]
        # # energy = [b]
        # a= self.Num[9]
        Pre=[]
        with torch.no_grad():
            # start = time.time()
            for i in range(n):
                # start = time.time()
                Snap_Test_Next = self.BECNN(Snap_init)
                # stop = time.time()
                # print(f"Training time: {stop - start}s")
                # Snap_init_np = self.Num[20* (i + 1) - 1]
                # Snap_init = torch.unsqueeze(torch.unsqueeze(torch.from_numpy(Snap_init_np), 0), 0).to(self.device)[:, :,
                #             1:, 0:-1]
                # Snap_Test_Next = Snap_init

                snap_full=self.full(Snap_Test_Next)
                Snap_init = Snap_Test_Next
                # a = np.max(np.abs(snap_full[0][0].cpu().detach().numpy()))
                # max.append(a)
                # Pre.append(snap_full[0][0].cpu().detach().numpy())
                predict = snap_full.mean(dim=0)[0].cpu().detach().numpy()
                # predict = snap_full[0][0].cpu().detach().numpy()
                Pre.append(predict)
                # a1,b1 = mass_engy(snap_full[0][0].cpu().detach().numpy(),self.hx,self.wx,self.epsilon)
                # mass.append(a1)
                # energy.append(b1)
                torch.cuda.empty_cache()
            np.save(f'data/deltaT{self.deltaT}_{self.height}_{self.width}_Ttrain=10_{888}_full.npy',
                    Pre)
        #     stop=time.time()
        # print(f"Training time: {stop - start}s")
        # One_step_predict = []
        # with torch.no_grad():
        #     for i in range(n):
        #         Snap_init_np = self.Num[5 * (i + 1) - 1]
        #         Snap_init = torch.unsqueeze(torch.unsqueeze(torch.from_numpy(Snap_init_np), 0), 0).to(self.device)[:, :,
        #                     1:, 0:-1]
        #         Snap_Test_Next = self.BECNN(Snap_init)
        #
        #         snap_full = self.full(Snap_Test_Next)
        #         next = snap_full[0][0].cpu().detach().numpy()
        #         One_step_predict.append(next)
        #         torch.cuda.empty_cache()
        # np.save(f'data/deltaT{0.05}_{self.height}_{self.width}_{n}_onestep.npy',
        #         One_step_predict)




        #     np.save(f'data/deltaT{self.deltaT}_{self.height}_{self.width}_2R_{self.mid_channels}_{666}_max.npy',
        #             max)
        #     np.save(f'data/deltaT{self.deltaT}_{self.height}_{self.width}_2R_{self.mid_channels}_{666}_mass.npy',
        #             mass)
        #     np.save(f'data/deltaT{self.deltaT}_{self.height}_{self.width}_2R_{self.mid_channels}_{666}_energy.npy',
        #             energy)
        #     pre = snap_full[0][0].cpu().detach().numpy()
        # plt.plot(np.linspace(0.1, 0.1*len(max), len(max)), max)
        # # plt.legend()
        # # plt.title("Maximum Value")
        # plt.xlabel('Time')
        # plt.ylabel('Maximum Value')
        # plt.show()
        # plt.plot(np.linspace(0.1, 0.1 * len(mass), len(mass)), mass)
        # # plt.legend()
        # # plt.title("Maximum Value")
        # plt.xlabel('Time')
        # plt.ylabel('Mass')
        # plt.show()
        # plt.plot(np.linspace(0.1, 0.1 * len(energy), len(energy)), energy)
        # # plt.legend()
        # # plt.title("Maximum Value")
        # plt.xlabel('Time')
        # plt.ylabel('Energy')
        # plt.show()

        # Snap_Numerical_Out0,Num = numerical_all(Snap_init0, self.domain, self.deltaT * n / 0.01, 0.01, self.epsilon,type='CN')
        # np.save(f'data/deltaT{1e-2}_{self.height}_{self.width}_{self.testBatch}_{1}_{88}_full.npy',
        #         Num)
        # # # # num = self.Numerical_Full[-1]
        # # num = Snap_Numerical_Out0
        # # # # Num =self.Numerical_Full
        # Numerical_all = Num
        # # # Pre=self.Pre
        # # # pre =Pre[-1]
        # RL2_error_all = []
        # for i in range(n):
        #     error = Num[10 * (i + 1) - 1] - Pre[i]
        #     RL2error = np.linalg.norm(error) / self.width
        #     RL2_error_all.append(RL2error)
        #
        # plt.plot(np.linspace(0.1, 0.1 * len(RL2_error_all), len(RL2_error_all)), RL2_error_all, label='0.1')
        # plt.legend()
        # plt.title("Error")
        # plt.show()

        Snap_Test_Sum = np.zeros((self.height + 1, self.width + 1))
        Numerical_1 = [0] * 5000*10
        for i in range(self.testBatch):
            Snap_init1 = torch.unsqueeze(Snap_init0[i], 0)
            Snap_Numerical_Out0, Numerical = numerical_all(Snap_init1, self.domain, self.deltaT*n/0.01, 0.01,
                                                           self.epsilon,
                                                           type='CN')
            # Snap_init1 = Snap_Numerical_Out0
            Numerical_1 = [a + b for a, b in zip(Numerical, Numerical_1)]

            Snap_Test_Sum = np.add(Snap_Test_Sum, Snap_Numerical_Out0)
        Snap_Numerical_Out = Snap_Test_Sum / self.testBatch
        Numerical_all = [a / self.testBatch for a in (Numerical_1)]
        # #
        np.save(f'data/deltaT{1e-2}_{self.height}_{self.width}_{888}_{20}_Numerical_all.npy',
                Numerical_all)



        Num=self.Num
        # Num=Numerical_all
        RL2_error_all = []
        for i in range(n):
            error = Num[5* (i + 1) - 1] - Pre[i]
            RL2error = np.linalg.norm(error) / self.width
            RL2_error_all.append(RL2error)

        plt.plot(np.linspace(0.05, 0.05* len(RL2_error_all), len(RL2_error_all)), RL2_error_all, label='0.05')
        plt.legend()
        plt.title("Error")
        plt.show()

        # def Average(lst):
        #     return sum(lst) / len(lst)
        #
        # One_1 = Average(RL2_error_all)
        # print(One_1)
        # Snap_Numerical_Out0 = numerical(Snap_init0, self.domain, self.deltaT * n / 0.1, 0.1, self.epsilon,
        #                                 type='CN')
        # np.save(f'data/deltaT{1e-3}_{self.height}_{self.width}_{self.deltaT * n * 1e+3}_{888}_full.npy',
        #         Snap_Numerical_Out0)
        # Snap_Numerical_Out0 = self.Numerical_Full


        # print(np.mean(snap_initial))
        # print(np.mean(snap_new_Num))
        # print(np.mean(snap_next))
        # pre =Pre[-1]
        # num =Num[-1]
        # # num = Snap_Numerical_Out
        # # error = np.abs(pre-num)
        # fig, ax = plt.subplots()
        # im = ax.imshow(pre, interpolation='bilinear', cmap=cm.RdYlGn)
        # plt.title(f'NN at T={n * self.deltaT}')
        # plt.colorbar(im)
        # plt.show()
        #
        # fig, ax1 = plt.subplots()
        # im1 = ax1.imshow(num, interpolation='bilinear', cmap=cm.RdYlGn)
        # plt.title(f'CN at T={n * self.deltaT}')
        # plt.colorbar(im1)
        # plt.show()
        #
        # fig, ax2 = plt.subplots()
        # im2 = ax2.imshow(error, interpolation='bilinear', cmap=cm.RdYlGn)
        # plt.title(f'Error at T={n * self.deltaT}')
        # plt.colorbar(im2)
        # plt.show()



if __name__ == '__main__':
    tester = Tester()
    tester.test(50)